import torch
import torch.nn.functional as F
from torch.optim import Adam
from torch_geometric.data import DataLoader
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch import tensor
import numpy as np
from utils import k_fold, num_graphs

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 定义了 train_causal_syn 函数，用于在合成数据集上训练因果图神经网络模型。
# 参数包括训练集 train_set、验证集 val_set、测试集 test_set、模型函数 model_func 和参数配置 args。
def train_causal_syn(train_set, val_set, test_set, model_func=None, args=None):

    # 使用 DataLoader 类将训练集、验证集和测试集转换为数据加载器。
    # 每个加载器的批量大小由 args.batch_size 定义，训练集被打乱，验证集和测试集不被打乱
    train_loader = DataLoader(train_set, args.batch_size, shuffle=True)
    val_loader = DataLoader(val_set, args.batch_size, shuffle=False)
    test_loader = DataLoader(test_set, args.batch_size, shuffle=False)

    # 如果特征维度 args.feature_dim 为 -1，则将其设置为最大度数 args.max_degree。
    if args.feature_dim == -1:
        args.feature_dim = args.max_degree

    # 使用模型函数 model_func 初始化因果图神经网络模型，并将其移动到设备上。
    # 使用 Adam 优化器初始化模型参数，学习率由 args.lr 定义。
    # 使用余弦退火学习率调度器初始化学习率调度器，设置最大迭代次数为 args.epochs。
    model = model_func(args.feature_dim, args.num_classes).to(device)
    optimizer = Adam(model.parameters(), lr=args.lr)
    lr_scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs, eta_min=args.min_lr, last_epoch=-1, verbose=False)

    # 初始化记录最佳验证准确率、更新的测试准确率（共同因果、因果、非因果）、以及更新的 epoch。
    best_val_acc, update_test_acc_co, update_test_acc_c, update_test_acc_o, update_epoch = 0, 0, 0, 0, 0

    # 开始迭代训练过程，遍历每个 epoch。
    for epoch in range(1, args.epochs + 1):

        # 调用 train_causal_epoch 函数执行一个 epoch 的训练，获取总损失、因果损失、因果效果损失、共同因果损失和训练准确率。
        train_loss, loss_c, loss_o, loss_co, train_acc_o = train_causal_epoch(model, optimizer, train_loader, device, args)

        # 使用验证集评估模型的性能，获取共同因果、因果和非因果模型的准确率。
        val_acc_co, val_acc_c, val_acc_o = eval_acc_causal(model, val_loader, device, args)
        # 使用测试集评估模型的性能，获取共同因果、因果和非因果模型的准确率。
        test_acc_co, test_acc_c, test_acc_o = eval_acc_causal(model, test_loader, device, args)
        # 更新学习率调度器的状态，进行学习率的更新。
        lr_scheduler.step()
        # 如果当前 epoch 的非因果模型在验证集上的准确率超过了历史最佳验证准确率，则更新记录的最佳验证准确率及其对应的测试准确率、共同因果准确率、因果准确率以及 epoch。
        if val_acc_o > best_val_acc:
            best_val_acc = val_acc_o
            update_test_acc_co = test_acc_co
            update_test_acc_c = test_acc_c
            update_test_acc_o = test_acc_o
            update_epoch = epoch

        # 打印每个 epoch 的训练损失、因果损失、因果效果损失、共同因果损失、训练准确率、验证准确率、测试准确率，以及记录的最佳测试准确率及其对应的 epoch 和学习率。
        print("BIAS:[{:.2f}] | Model:[{}] Epoch:[{}/{}] Loss:[{:.4f}={:.4f}+{:.4f}+{:.4f}] Train:[{:.2f}] val:[{:.2f}] Test:[{:.2f}] | Update Test:[co:{:.2f},c:{:.2f},o:{:.2f}] at Epoch:[{}] | lr:{:.6f}"
                .format(args.bias,
                        args.model,
                        epoch, 
                        args.epochs,
                        train_loss,
                        loss_c,
                        loss_o,
                        loss_co,
                        train_acc_o * 100, 
                        val_acc_o * 100,
                        test_acc_o * 100, 
                        update_test_acc_co * 100,
                        update_test_acc_c * 100,  
                        update_test_acc_o * 100, 
                        update_epoch,
                        optimizer.param_groups[0]['lr']))
    # 打印最终结果，包括验证集准确率以及更新的测试准确率（共同因果、因果、非因果）及其对应的 epoch。
    print("syd: BIAS:[{:.2f}] | Val acc:[{:.2f}] Test acc:[co:{:.2f},c:{:.2f},o:{:.2f}] at epoch:[{}]"
        .format(args.bias,
                val_acc_o * 100,
                update_test_acc_co * 100,
                update_test_acc_c * 100,  
                update_test_acc_o * 100, 
                update_epoch))

# 定义了 train_causal_real 函数，用于在真实数据集上训练因果图神经网络模型。
# 参数包括数据集 dataset、模型函数 model_func 和参数配置 args。
def train_causal_real(dataset=None, model_func=None, args=None):

    # 初始化记录训练准确率、测试准确率（共同因果、因果、非因果）的列表，并计算随机猜测的准确率。
    train_accs, test_accs, test_accs_c, test_accs_o = [], [], [], []
    random_guess = 1.0 / dataset.num_classes
    # 对数据集进行 k 折交叉验证，其中 k_fold() 函数返回了训练集、验证集和测试集的索引。
    for fold, (train_idx, test_idx, val_idx) in enumerate(zip(*k_fold(dataset, args.folds, args.epoch_select))):
        # 对每个折（fold）进行训练和测试准备。初始化记录最佳测试准确率、最佳 epoch、最佳测试因果准确率、最佳测试非因果准确率，并分别加载训练集和测试集。
        best_test_acc, best_epoch, best_test_acc_c, best_test_acc_o = 0, 0, 0, 0
        train_dataset = dataset[train_idx]
        test_dataset = dataset[test_idx]
        train_loader = DataLoader(train_dataset, args.batch_size, shuffle=True)
        test_loader = DataLoader(test_dataset, args.batch_size, shuffle=False)

        # 使用模型函数 model_func 初始化因果图神经网络模型，并将其移动到设备上。使用 Adam 优化器初始化模型参数，学习率由 args.lr 定义，同时应用权重衰减。
        model = model_func(dataset.num_features, dataset.num_classes).to(device)
        optimizer = Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

        # 开始迭代训练过程，遍历每个 epoch。
        for epoch in range(1, args.epochs + 1):
            # 调用 train_causal_epoch() 函数执行一个 epoch 的训练，获取总损失、因果损失、因果效果损失、共同因果损失和训练准确率。然后使用测试集评估模型的性能，获取测试准确率、测试因果准确率和测试非因果准确率。
            train_loss, loss_c, loss_o, loss_co, train_acc = train_causal_epoch(model, optimizer, train_loader, device, args)
            test_acc, test_acc_c, test_acc_o = eval_acc_causal(model, test_loader, device, args)

            # 将每个 epoch 的训练准确率、测试准确率（共同因果、因果、非因果）添加到对应的列表中。
            train_accs.append(train_acc)
            test_accs.append(test_acc)
            test_accs_c.append(test_acc_c)
            test_accs_o.append(test_acc_o)

            # 如果当前 epoch 的测试准确率超过了历史最佳测试准确率，则更新记录的最佳测试准确率及其对应的 epoch、最佳测试因果准确率和最佳测试非因果准确率。
            if test_acc > best_test_acc:
                best_test_acc = test_acc
                best_epoch = epoch
                best_test_acc_c = test_acc_c
                best_test_acc_o = test_acc_o

            # 打印每个 epoch 的训练损失、因果损失、因果效果损失、共同因果损失、训练准确率、测试准确率、测试非因果准确率、测试因果准确率，以及记录的最佳测试准确率及其对应的 epoch 和随机猜测准确率。
            print("Causal | dataset:[{}] fold:[{}] | Epoch:[{}/{}] Loss:[{:.4f}={:.4f}+{:.4f}+{:.4f}] Train:[{:.4f}] Test:[{:.2f}] Test_o:[{:.2f}] Test_c:[{:.2f}] (RG:{:.2f}) | Best Test:[{:.2f}] at Epoch:[{}] | Test_o:[{:.2f}] Test_c:[{:.2f}]"
                    .format(args.dataset,
                            fold,
                            epoch, args.epochs,
                            train_loss, loss_c, loss_o, loss_co,
                            train_acc * 100,  
                            test_acc * 100, 
                            test_acc_o * 100,
                            test_acc_c * 100, 
                            random_guess*  100,
                            best_test_acc * 100, 
                            best_epoch,
                            best_test_acc_o * 100,
                            best_test_acc_c * 100))

        # 打印每个 fold 的最终结果，包括最佳测试准确率及其对应的 epoch、最佳测试非因果准确率、随机猜测准确率。
        print("syd: Causal fold:[{}] | Dataset:[{}] Model:[{}] | Best Test:[{:.2f}] at epoch [{}] | Test_o:[{:.2f}] Test_c:[{:.2f}] (RG:{:.2f})"
                .format(fold,
                        args.dataset,
                        args.model,
                        best_test_acc * 100, 
                        best_epoch,
                        best_test_acc_o * 100,
                        best_test_acc_c * 100,
                        random_guess*  100))

    # 将列表中的训练准确率、测试准确率（共同因果、因果、非因果）转换为 PyTorch 张量。
    train_acc, test_acc, test_acc_c, test_acc_o = tensor(train_accs), tensor(test_accs), tensor(test_accs_c), tensor(test_accs_o)
    # 根据折数和 epoch 数将张量重塑为二维形状，其中每行表示一个折的训练或测试结果。
    train_acc = train_acc.view(args.folds, args.epochs)
    test_acc = test_acc.view(args.folds, args.epochs)
    test_acc_c = test_acc_c.view(args.folds, args.epochs)
    test_acc_o = test_acc_o.view(args.folds, args.epochs)

    # 计算测试准确率在所有折上的均值，并找到具有最高均值的 epoch。将选定的 epoch 扩展为与折数相同的张量。
    _, selected_epoch = test_acc.mean(dim=0).max(dim=0)
    selected_epoch = selected_epoch.repeat(args.folds)
    
    _, selected_epoch2 = test_acc_o.mean(dim=0).max(dim=0)
    selected_epoch2 = selected_epoch2.repeat(args.folds)

    # 从测试准确率张量中选取在每个折上具有最高均值的测试准确率，并将其提取为一维张量。
    test_acc = test_acc[torch.arange(args.folds, dtype=torch.long), selected_epoch]
    test_acc_c = test_acc_c[torch.arange(args.folds, dtype=torch.long), selected_epoch]
    test_acc_o = test_acc_o[torch.arange(args.folds, dtype=torch.long), selected_epoch2]

    # 计算训练准确率的平均值、测试准确率（共同因果、因果、非因果）的平均值和标准差。
    train_acc_mean = train_acc[:, -1].mean().item()
    test_acc_mean = test_acc.mean().item()
    test_acc_std = test_acc.std().item()
    test_acc_c_mean = test_acc_c.mean().item()
    test_acc_c_std = test_acc_c.std().item()
    test_acc_o_mean = test_acc_o.mean().item()
    test_acc_o_std = test_acc_o.std().item()

    # 打印最终的实验结果，包括测试准确率的均值和标准差，以及共同因果、因果、非因果测试准确率的均值和标准差，同时打印了模型的设置参数。
    print("=" * 150)
    print('sydall Final: Causal | Dataset:[{}] Model:[{}] seed:[{}]| Test Acc: {:.2f}±{:.2f} | OTest: {:.2f}±{:.2f}, CTest: {:.2f}±{:.2f} (RG:{:.2f}) | [Settings] co:{},c:{},o:{},harf:{},dim:{},fc:{}'
         .format(args.dataset,
                 args.model,
                 args.seed,
                 test_acc_mean * 100, 
                 test_acc_std * 100,
                 test_acc_o_mean * 100, 
                 test_acc_o_std * 100,
                 test_acc_c_mean * 100, 
                 test_acc_c_std * 100,
                 random_guess*  100,
                 args.co,
                 args.c,
                 args.o,
                 args.harf_hidden,
                 args.hidden,
                 args.fc_num))
    print("=" * 150)

# 此函数用于训练因果图神经网络的一个 epoch。
def train_causal_epoch(model, optimizer, loader, device, args):
    # 将模型设置为训练模式，并初始化总损失和总损失的各个组成部分（共同因果、因果、非因果），以及正确分类的非因果节点数。
    model.train()
    total_loss = 0
    total_loss_c = 0
    total_loss_o = 0
    total_loss_co = 0
    correct_o = 0
    # 迭代数据加载器中的每个批次。
    for it, data in enumerate(loader):
        # 将优化器的梯度置零，并将数据移动到适当的设备上（GPU或CPU）。
        optimizer.zero_grad()
        data = data.to(device)
        # 将目标标签转换为独热编码形式，并通过模型获取共同因果、因果和非因果节点的预测概率。还创建了一个均匀分布的目标概率张量。
        one_hot_target = data.y.view(-1)
        c_logs, o_logs, co_logs = model(data, eval_random=args.with_random)
        uniform_target = torch.ones_like(c_logs, dtype=torch.float).to(device) / model.num_classes
        # 分别计算共同因果、因果和非因果节点的 KL 散度损失、负对数似然损失，并将它们的加权和作为总损失。
        c_loss = F.kl_div(c_logs, uniform_target, reduction='batchmean')
        o_loss = F.nll_loss(o_logs, one_hot_target)
        co_loss = F.nll_loss(co_logs, one_hot_target)
        loss = args.c * c_loss + args.o * o_loss + args.co * co_loss
        # 找到非因果节点的预测值，并更新正确分类的非因果节点数。然后根据总损失计算梯度，并执行反向传播。
        pred_o = o_logs.max(1)[1]
        correct_o += pred_o.eq(data.y.view(-1)).sum().item()
        loss.backward()
        # 计算并累积总损失和各个组成部分的损失。随后执行一步优化器更新模型参数。
        total_loss += loss.item() * num_graphs(data)
        total_loss_c += c_loss.item() * num_graphs(data)
        total_loss_o += o_loss.item() * num_graphs(data)
        total_loss_co += co_loss.item() * num_graphs(data)
        optimizer.step()
    # 计算损失的平均值，并将非因果节点的正确分类比例转换为百分比。
    num = len(loader.dataset)
    total_loss = total_loss / num
    total_loss_c = total_loss_c / num
    total_loss_o = total_loss_o / num
    total_loss_co = total_loss_co / num
    correct_o = correct_o / num
    # 返回共同因果、因果和非因果节点的平均损失以及正确分类的非因果节点的比例。
    return total_loss, total_loss_c, total_loss_o, total_loss_co, correct_o

# 该函数用于评估因果图神经网络的性能。
def eval_acc_causal(model, loader, device, args):
    # 将模型设置为评估模式，并根据参数 args 确定是否要在评估时启用随机评估模式。
    model.eval()
    eval_random = args.eval_random
    # 初始化共同因果、因果和非因果节点的正确分类计数。
    correct = 0
    correct_c = 0
    correct_o = 0
    # 迭代数据加载器中的每个批次。
    for data in loader:
        # 将数据移动到适当的设备上（GPU或CPU）。
        data = data.to(device)
        # 通过模型获取共同因果、因果和非因果节点的预测概率，并找到预测概率最高的类别。
        with torch.no_grad():
            c_logs, o_logs, co_logs = model(data, eval_random=eval_random)
            pred = co_logs.max(1)[1]
            pred_c = c_logs.max(1)[1] 
            pred_o = o_logs.max(1)[1]
        # 更新共同因果、因果和非因果节点的正确分类计数。
        correct += pred.eq(data.y.view(-1)).sum().item()
        correct_c += pred_c.eq(data.y.view(-1)).sum().item()
        correct_o += pred_o.eq(data.y.view(-1)).sum().item()
    # 计算共同因果、因果和非因果节点的分类准确率。
    acc_co = correct / len(loader.dataset)
    acc_c = correct_c / len(loader.dataset)
    acc_o = correct_o / len(loader.dataset)
    # 返回共同因果、因果和非因果节点的分类准确率。
    return acc_co, acc_c, acc_o
